#include "../DWMRICore/Fiber.h"
#include "../DWMRICore/DTI.h"
#include "../DWMRICore/HOT.h"
#include "../DWMRICore/Scalar.h"

#include <omp.h>

#define _USE_MATH_DEFINES
#include <math.h>
#include <time.h>

#include <teem/ten.h>

#include <string>

using namespace std;

char *pDTIPathname = NULL;					/* input DTI file pathname */
char *pHOTPathname = NULL;					/* input HOT file pathname */
char *pHOT6Pathname = NULL;					/* input 6th HOT file pathname */
char *pQBallPathname = NULL;				/* input QBall file pathname */
char *pBrainMASKPathname = NULL;			/* input brain mask file pathname */
char *pWMMASKPathname = NULL;				/* input white matter mask file pathname */
char *pModelSelectionPathname = NULL;		/* input model selection file pathname */
char *pFiberPathname = NULL;
char *pSinglePathname = NULL;
char *pYMASKPathname = NULL;

float stepSize = 0.25f;
float maxLength = 15.0f;
float faThreshold = 0.2f;
float angleThreshold = 0.866f;
int scale = 1;

/* MLS fiber tracking parameters */
float MLSSize = 3.0f;
CPoint3F *pSamplePoints = NULL;
float *pSamplePointsWeight = NULL;
float *matG = NULL;

CDTI *pDTI = NULL;					/* DTI */
CHOT *pHOT = NULL;					/* HOT */
CHOT *pQBall = NULL;
CHOT *pHOT6 = NULL;					/* 6th order tensor */
CScalar *pBrainMASK = NULL;			/* brain mask */
CScalar *pWMMASK = NULL;			/* white matter mask */
CScalar *pModelSelection = NULL;	/* model selection, DTI or HOT */
CScalar *pSingleMASK = NULL;		/* single orientation mask */
CScalar *pYMASK = NULL;				/* Y mask */

/* second and third orientations */
CVector3F *pSecondOrientation = NULL;
CVector3F *pThirdOrientation = NULL;

bool isFirstFlag = true;
bool isTrackingFlag = false;
int startY = 584;
int endY = 656;

void PrintUsage()
{
	printf("Usage: ");
	printf("FiberTracking -dti <DTI> -hot4 <4th HOT> -hot6 <6th HOT> -qball <QBall> -y_mask <Y MASK> -single <Single MASK> -brain_mask <Brain MASK> -wm_mask <White Matter MASK> -model_selection <Model Selection> -fiber <Fiber> [-stepsize <stepsize> -maxlength <maxlength> -fa_thresold <fa_threshold> -angle_threshold <fa_threshold> -scale <scale>]\n");
	printf("\n");
	printf("-vec <VectorField> = input VectorField nrrd file\n");
	printf("-fiber <Fiber> = output Fiber configure file\n");
	printf("-stepsize <stepsize> = step size\n");
	printf("-maxlength <maxlength> = maximum fiber length\n");
	printf("-fa_threshold <fa_thresold> = FA threshold used to differ white matter and grey matter\n");
	printf("-angle_threshold <angle_thresold> = angle threshold between two directions, in cosine\n");
	printf("-scale <scale> = scale of the result\n");
}

int ParseArguments(const int argc, char *argv[])
{
	for (int i = 1; i < argc; ++i) {
		if (strcmp(argv[i], "-dti") == 0) {
			pDTIPathname = argv[++i];
		} else if (strcmp(argv[i], "-hot4") == 0) {
			pHOTPathname = argv[++i];
		} else if (strcmp(argv[i], "-hot6") == 0) {
			pHOT6Pathname = argv[++i];
		} else if (strcmp(argv[i], "-qball") == 0) {
			pQBallPathname = argv[++i];
		} else if (strcmp(argv[i], "-single") == 0) {
			pSinglePathname = argv[++i];
		} else if (strcmp(argv[i], "-y_mask") == 0) {
			pYMASKPathname = argv[++i];
		} else if (strcmp(argv[i], "-brain_mask") == 0) {
			pBrainMASKPathname = argv[++i];
		} else if (strcmp(argv[i], "-wm_mask") == 0) {
			pWMMASKPathname = argv[++i];
		} else if (strcmp(argv[i], "-model_selection") == 0) {
			pModelSelectionPathname = argv[++i];
		} else if (strcmp(argv[i], "-fiber") == 0) {
			pFiberPathname = argv[++i];
		} else if (strcmp(argv[i], "-stepsize") == 0) {
			stepSize = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-maxlength") == 0) {
			maxLength = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-fa_threshold") == 0) {
			faThreshold = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-angle_threshold") == 0) {
			angleThreshold = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-scale") == 0) {
			scale = (uint)atoi(argv[++i]);
		} else {
            printf("unknown %s\n", argv[i]);
			return 1;		// error
		}
	}

	if (pDTIPathname == NULL || pHOTPathname == NULL || pQBallPathname == NULL || 
		pBrainMASKPathname == NULL || pWMMASKPathname == NULL || pFiberPathname == NULL)
		return -2;

	return 0;
}

void CreateMLSSamplePoints()
{
	SafeDeleteArray(pSamplePoints);
	SafeDeleteArray(pSamplePointsWeight);
	pSamplePoints = new CPoint3F[SAMPLE_POINT_NUM];
	pSamplePointsWeight = new float[SAMPLE_POINT_NUM];

	pSamplePoints[0].m_x = 0.0f;
	pSamplePoints[0].m_y = 0.0f;
	pSamplePoints[0].m_z = 0.0f;
	pSamplePointsWeight[0] = 1.0f;

	clock_t seed = clock();
	srand(seed);

	for (int i = 1; i < SAMPLE_POINT_NUM; ++i) {
		bool flag = false;
		CPoint3F pt;
		while(!flag) {
			float a = (float)rand() / (float)RAND_MAX * 2.0f * (float)M_PI;
			float b = (float)rand() / (float)RAND_MAX * (float)M_PI;
			float r = (float)rand() / (float)RAND_MAX;

			pt.m_x = r * sinf(a) * sinf(b);
			pt.m_y = r * sinf(a) * cosf(b);
			pt.m_z = r * cosf(a);

			flag = true;
			for (int j = 0; j < i; ++j) {
				float dis = Distance(pSamplePoints[j], pt);
				if (dis < 0.33f) {
					flag = false;
					break;
				}
			}
		}

		pSamplePoints[i] = pt;
		pSamplePointsWeight[i] = InnerProduct(pSamplePoints[i], pSamplePoints[i]) * 4.0f;
		pSamplePointsWeight[i] = expf(pSamplePointsWeight[i] / -2.5f);
	}
}

void CreateMatG()
{
	int order = pQBall->GetOrder();
	int n = pQBall->GetN();
	
	SafeDeleteArray(matG);
	matG = new float[SAMPLE_POINT_NUM*n];

	for (int i = 0; i < SAMPLE_POINT_NUM; ++i) {
		int k = 0;
		CVector3F v = pSamplePoints[i];
		for (int a = 0; a <= order; ++a) {
			for (int b = 0; b <= order-a; ++b) {
				matG[i*n+k] = pow(v.m_x, a) * pow(v.m_y, b) * pow(v.m_z, (order-a-b));
				k++;
			}
		}
	}
}


void ConvertQBall(float *src, float *dst)
{
	float factor[15] = {1, 4, 6, 4, 1, 4, 12, 12, 4, 6, 12, 6, 4, 4, 1};
	for (int i = 0; i < 15; ++i) {
		dst[i] = src[15-i-1] * factor[i];
	}
}

/************************************************************/
/* filter the DTI data by MLS								*/
/************************************************************/
void DTIMLSFilter(const CPoint3F pos, float *filteredDTI)
{
	float dti[7];
	pDTI->GetDTI(pos, dti);
	float m[9];
	TEN_T2M(m, dti);

	CPoint3F samples[SAMPLE_POINT_NUM];
	float length[SAMPLE_POINT_NUM];
	float max_length = 0.0f;
	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		samples[i].m_x = pSamplePoints[i].m_x * m[0] + pSamplePoints[i].m_y * m[1] + pSamplePoints[i].m_z * m[2];
		samples[i].m_y = pSamplePoints[i].m_x * m[3] + pSamplePoints[i].m_y * m[4] + pSamplePoints[i].m_z * m[5];
		samples[i].m_z = pSamplePoints[i].m_x * m[6] + pSamplePoints[i].m_y * m[7] + pSamplePoints[i].m_z * m[8];
		length[i] = sqrtf(InnerProduct(samples[i], samples[i]));
		if (length[i] > max_length)
			max_length = length[i];
	}

	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		samples[i].m_x = samples[i].m_x / max_length * MLSSize + pos.m_x;
		samples[i].m_y = samples[i].m_y / max_length * MLSSize + pos.m_y;
		samples[i].m_z = samples[i].m_z / max_length * MLSSize + pos.m_z;
	}

	memset(filteredDTI, 0, sizeof(float)*7);
	float temp_DTI[7];
	float weight = 0.0f;
	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		if (pBrainMASK->GetScalar(samples[i]) < 0.99f) {
			continue;
		}
		//if (pWMMASK->GetScalar(samples[i]) > 0.4f) {
		//	continue;
		//}
		pDTI->GetDTI(samples[i], temp_DTI);

		for (uint j = 0; j < 7; ++j) {
			filteredDTI[j] += pSamplePointsWeight[i] * temp_DTI[j];
		}
		weight += pSamplePointsWeight[i];
	}

	for (uint j = 0; j < 7; ++j) {
		filteredDTI[j] = filteredDTI[j] / weight;
	}
}

//void DTIMLSFilter(const CPoint3F pos, float *out_DTI, float *in_DTI)
//{
//	float m[9];
//	TEN_T2M(m, in_DTI);
//
//	CPoint3F samples[SAMPLE_POINT_NUM];
//	float length[SAMPLE_POINT_NUM];
//	float max_length = 0.0f;
//	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
//		samples[i].m_x = pSamplePoints[i].m_x * m[0] + pSamplePoints[i].m_y * m[1] + pSamplePoints[i].m_z * m[2];
//		samples[i].m_y = pSamplePoints[i].m_x * m[3] + pSamplePoints[i].m_y * m[4] + pSamplePoints[i].m_z * m[5];
//		samples[i].m_z = pSamplePoints[i].m_x * m[6] + pSamplePoints[i].m_y * m[7] + pSamplePoints[i].m_z * m[8];
//		length[i] = sqrtf(InnerProduct(samples[i], samples[i]));
//		if (length[i] > max_length)
//			max_length = length[i];
//	}
//
//	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
//		samples[i].m_x = samples[i].m_x / max_length * MLSSize + pos.m_x;
//		samples[i].m_y = samples[i].m_y / max_length * MLSSize + pos.m_y;
//		samples[i].m_z = samples[i].m_z / max_length * MLSSize + pos.m_z;
//	}
//
//	memset(out_DTI, 0, sizeof(float)*7);
//	float temp_DTI[7];
//	float weight = 0.0f;
//	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
//		if (pBrainMASK->GetScalar(samples[i]) < 0.5f) {
//			continue;
//		}
//		pDTI->GetDTI(samples[i], temp_DTI);
//
//		for (uint j = 0; j < 7; ++j) {
//			out_DTI[j] += pSamplePointsWeight[i] * temp_DTI[j];
//		}
//		weight += pSamplePointsWeight[i];
//	}
//
//	for (uint j = 0; j < 7; ++j) {
//		out_DTI[j] = out_DTI[j] / weight;
//	}
//}

/************************************************************/
/* filter the HOT data by MLS								*/
/************************************************************/
void HOT6MLSFilter(const CPoint3F pos, float *filteredHOT)
{
	uint n = pQBall->GetN();

	float *hot = new float[n];
	float *qball = new float[n];
	pQBall->GetHOT(pos, hot);
	ConvertQBall(hot, qball);

	CPoint3F samples[SAMPLE_POINT_NUM];
	float length[SAMPLE_POINT_NUM];
	float max_length = 0.0f;
	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		length[i] = 0.0f;
		for (uint j = 0; j < n; ++j) {
			length[i] += matG[i*n+j] * qball[j];
		}
		if (length[i] < 0.0f)
			length[i] = 0.0f;
		if (length[i] > max_length)
			max_length = length[i];
	}

	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		samples[i].m_x = pSamplePoints[i].m_x / max_length * MLSSize + pos.m_x;
		samples[i].m_y = pSamplePoints[i].m_y / max_length * MLSSize + pos.m_y;
		samples[i].m_z = pSamplePoints[i].m_z / max_length * MLSSize + pos.m_z;
	}

	n = pHOT6->GetN();
	memset(filteredHOT, 0, sizeof(float)*n);
	float *temp_HOT = new float[n];
	float weight = 0.0f;
	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		if (pBrainMASK->GetScalar(samples[i]) < 0.5f) {
			continue;
		}
		//if (pWMMASK->GetScalar(samples[i]) > 0.4f) {
		//	continue;
		//}
		pHOT6->GetHOT(samples[i], temp_HOT);

		for (uint j = 0; j < n; ++j) {
			filteredHOT[j] += pSamplePointsWeight[i] * temp_HOT[j];
		}
		weight += pSamplePointsWeight[i];
	}

	for (uint j = 0; j < n; ++j) {
		filteredHOT[j] = filteredHOT[j] / weight;
	}

	SafeDeleteArray(temp_HOT);
	SafeDeleteArray(hot);
	SafeDeleteArray(qball);
}

void HOTMLSFilter(const CPoint3F pos, float *filteredHOT)
{
	uint n = pQBall->GetN();

	float *hot = new float[n];
	float *qball = new float[n];
	pQBall->GetHOT(pos, hot);
	ConvertQBall(hot, qball);

	CPoint3F samples[SAMPLE_POINT_NUM];
	float length[SAMPLE_POINT_NUM];
	float max_length = 0.0f;
	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		length[i] = 0.0f;
		for (uint j = 0; j < n; ++j) {
			length[i] += matG[i*n+j] * qball[j];
		}
		if (length[i] < 0.0f)
			length[i] = 0.0f;
		if (length[i] > max_length)
			max_length = length[i];
	}

	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		samples[i].m_x = pSamplePoints[i].m_x / max_length * MLSSize + pos.m_x;
		samples[i].m_y = pSamplePoints[i].m_y / max_length * MLSSize + pos.m_y;
		samples[i].m_z = pSamplePoints[i].m_z / max_length * MLSSize + pos.m_z;
	}

	n = pHOT->GetN();
	memset(filteredHOT, 0, sizeof(float)*n);
	float *temp_HOT = new float[n];
	float weight = 0.0f;
	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		if (pBrainMASK->GetScalar(samples[i]) < 0.5f) {
			continue;
		}
		//if (pWMMASK->GetScalar(samples[i]) > 0.4f) {
		//	continue;
		//}
		pHOT->GetHOT(samples[i], temp_HOT);

		for (uint j = 0; j < n; ++j) {
			filteredHOT[j] += pSamplePointsWeight[i] * temp_HOT[j];
		}
		weight += pSamplePointsWeight[i];
	}

	for (uint j = 0; j < n; ++j) {
		filteredHOT[j] = filteredHOT[j] / weight;
	}

	SafeDeleteArray(temp_HOT);
	SafeDeleteArray(hot);
	SafeDeleteArray(qball);
}
//
//void HOTMLSFilter(const CPoint3F pos, float *out_HOT, float *in_DTI)
//{
//	float m[9];
//	TEN_T2M(m, in_DTI);
//
//	CPoint3F samples[SAMPLE_POINT_NUM];
//	float length[SAMPLE_POINT_NUM];
//	float max_length = 0.0f;
//	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
//		samples[i].m_x = pSamplePoints[i].m_x * m[0] + pSamplePoints[i].m_y * m[1] + pSamplePoints[i].m_z * m[2];
//		samples[i].m_y = pSamplePoints[i].m_x * m[3] + pSamplePoints[i].m_y * m[4] + pSamplePoints[i].m_z * m[5];
//		samples[i].m_z = pSamplePoints[i].m_x * m[6] + pSamplePoints[i].m_y * m[7] + pSamplePoints[i].m_z * m[8];
//		length[i] = sqrtf(InnerProduct(samples[i], samples[i]));
//		if (length[i] > max_length)
//			max_length = length[i];
//	}
//
//	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
//		samples[i].m_x = samples[i].m_x / max_length * MLSSize + pos.m_x;
//		samples[i].m_y = samples[i].m_y / max_length * MLSSize + pos.m_y;
//		samples[i].m_z = samples[i].m_z / max_length * MLSSize + pos.m_z;
//	}
//
//	uint n = pHOT->GetN();
//	memset(out_HOT, 0, sizeof(float)*n);
//	float *temp_HOT = new float[n];
//	float weight = 0.0f;
//	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
//		if (pBrainMASK->GetScalar(samples[i]) < 0.5f) {
//			continue;
//		}
//		pHOT->GetHOT(samples[i], temp_HOT);
//
//		for (uint j = 0; j < n; ++j) {
//			out_HOT[j] += pSamplePointsWeight[i] * temp_HOT[j];
//		}
//		weight += pSamplePointsWeight[i];
//	}
//	delete[] temp_HOT;
//
//	for (uint j = 0; j < n; ++j) {
//		out_HOT[j] = out_HOT[j] / weight;
//	}
//}

int ModelSelection(const CPoint3F pos)
{
	float x = pos.m_x;
	float y = pos.m_y;
	float z = pos.m_z;

	if (pBrainMASK->GetScalar(x, y, z) < 0.5f) {
		return 0;			// not in the brain
	} else {
		if (pModelSelection == NULL) {
			return 2;		// HOT
		} else {
			float model = pModelSelection->GetScalar(x, y, z);
			if (model > 0.5f)
				return 1;		// DTI
			else
				return 2;		// HOT
		}
	}
}

void AddOrientation(const CPoint3F pos, const CVector3F dir)
{
	if (isTrackingFlag == true)
		return;

	int w = pHOT->GetWidth();
	int h = pHOT->GetHeight();
	int d = pHOT->GetDepth();

	int sw = w * scale;
	int sh = h * scale;
	int sd = d * scale;
	int dy = endY - startY;

	int x = int(pos.m_x * scale + 0.5f);
	int y = int(pos.m_y * scale + 0.5f);
	int z = int(pos.m_z * scale + 0.5f);

	if (y < startY || y >= endY)
		return;

	int index = (z * dy + (y-startY)) * sw + x;

	CPoint3F pp;
	pp.m_x = (float)x / (float)scale;
	pp.m_y = (float)y / (float)scale;
	pp.m_z = (float)z / (float)scale;

	int model = ModelSelection(pp);
	if (model != 2) {
		return;
	}

	uint n = pHOT->GetN();
	float *filteredHOT = new float[n];
	HOTMLSFilter(pp, filteredHOT);

	CVector3F dirs[3];
	float weights[3] = {0.0f, 0.0f, 0.0f};
	int count = pHOT->HOT2Directions(filteredHOT, dirs, weights);
	delete[] filteredHOT;

	if (count == 1 || fabs(InnerProduct(dir, dirs[0])) > 0.866f) {
		return;
	}

	if (count > 1 && fabs(InnerProduct(dir, dirs[1])) > 0.866f) {
		pSecondOrientation[index] = dirs[1];
		return;
	}

	if (count > 2 && fabs(InnerProduct(dir, dirs[2])) > 0.866f) {
		pThirdOrientation[index] = dirs[2];
		return;
	}
}

int GenerateSeeds(const CPoint3F pos, Seed **seeds)
{
	int model = ModelSelection(pos);

	if (pos.m_z < 55.0f || pos.m_z > 79.9f || pos.m_x < 30.0f || pos.m_x > 59.9f) {
		// not in the interesting region
		*seeds = new Seed[1];
		(*seeds)[0].pos = pos;
		(*seeds)[0].dir = CVector3F(0.0f, 0.0f, 0.0f);
		return 1;
	}

	float filteredDTI[7];
	DTIMLSFilter(pos, filteredDTI);
	float fa = pDTI->DTI2FA(filteredDTI);

	if (model == 0 || (fa < faThreshold && pWMMASK->GetScalar(pos) < 0.5f) || 
		(fa < 0.1f)) {
		// not in the white matter
		*seeds = new Seed[1];
		(*seeds)[0].pos = pos;
		(*seeds)[0].dir = CVector3F(0.0f, 0.0f, 0.0f);
		return 1;
	} else {
		if (model == 1) {
			// DTI
			CVector3F dir;
			pDTI->DTI2Direction(filteredDTI, &dir);
			*seeds = new Seed[1];
			(*seeds)[0].pos = pos;
			(*seeds)[0].dir = dir;
			return 1;
		} else {
			// HOT
			uint n = pHOT->GetN();
			float *filteredHOT = new float[n];
			HOTMLSFilter(pos, filteredHOT);

			CVector3F dirs[3];
			float weights[3] = {0.0f, 0.0f, 0.0f};
			int count = pHOT->HOT2Directions(filteredHOT, dirs, weights);
			delete[] filteredHOT;

			int number = 1;
			for (int i = 1; i < count; ++i) {
				if (weights[i] > weights[0] * 0.9f) {
					number++;
					//AddOrientation(pos, dirs[i]);
				}
			}

			*seeds = new Seed[number];
			for (int i = 0; i < number; ++i) {
				(*seeds)[i].pos = pos;
				(*seeds)[i].dir = dirs[i];
			}

			return number;
		}
	}
}

int GenerateSeedsNew(const CPoint3F pos, Seed **seeds)
{
	int model = ModelSelection(pos);

	//if (pos.m_z < 66.0f || pos.m_z > 74.9f || pos.m_x < 37.0f || pos.m_x > 43.9f) {
	if (pos.m_z < 55.0f || pos.m_z > 79.9f || pos.m_x < 30.0f || pos.m_x > 59.9f) {
		// not in the interesting region
		*seeds = new Seed[1];
		(*seeds)[0].pos = pos;
		(*seeds)[0].dir = CVector3F(0.0f, 0.0f, 0.0f);
		return 1;
	}

	float filteredDTI[7];
	DTIMLSFilter(pos, filteredDTI);
	float fa = pDTI->DTI2FA(filteredDTI);

	if (model == 0 || (fa < faThreshold && pWMMASK->GetScalar(pos) < 0.5f) || 
		(fa < 0.1f)) {
	//if (model == 0 || (fa < faThreshold || pWMMASK->GetScalar(pos) > 0.4f)) {
		// not in the white matter
		*seeds = new Seed[1];
		(*seeds)[0].pos = pos;
		(*seeds)[0].dir = CVector3F(0.0f, 0.0f, 0.0f);
		return 1;
	} else {
		if (model == 1) {
			// DTI
			CVector3F dir;
			pDTI->DTI2Direction(filteredDTI, &dir);
			*seeds = new Seed[1];
			(*seeds)[0].pos = pos;
			(*seeds)[0].dir = dir;
			return 1;
		} else {
			// HOT
			uint n = pHOT->GetN();
			float *filteredHOT = new float[n];
			HOTMLSFilter(pos, filteredHOT);

			CVector3F dirs[3];
			float weights[3] = {0.0f, 0.0f, 0.0f};
			int count = pHOT->HOT2Directions(filteredHOT, dirs, weights);
			delete[] filteredHOT;

			int w = pHOT->GetWidth();
			int h = pHOT->GetHeight();
			int d = pHOT->GetDepth();

			int sw = w * scale;
			int sh = h * scale;
			int sd = d * scale;
			int dy = endY - startY;

			int x = int(pos.m_x * scale + 0.5f);
			int y = int(pos.m_y * scale  + 0.5f);
			int z = int(pos.m_z * scale  + 0.5f);
			int index = (z * dy + (y-startY)) * sw + x;

			int number;
			if(isFirstFlag) {
				number = 1;
			} else {
				number = 0;
			}
			
			for (int i = 1; i < count; ++i) {
				if (weights[i] > weights[0] * 0.8f || 
					fabs(InnerProduct(pSecondOrientation[index], dirs[i])) > 0.866f || 
					fabs(InnerProduct(pThirdOrientation[index], dirs[i])) > 0.866f) {
					number++;
					//AddOrientation(pos, dirs[i]);
				}
			}

			if (number == 0)
				return 0;
			
			if (isFirstFlag && pSingleMASK != NULL && pSingleMASK->GetScalar(pos) > 0.5f) {
				number = 1;
			}

			*seeds = new Seed[number];
			int i;
			if(isFirstFlag) {
				(*seeds)[0].pos = pos;
				(*seeds)[0].dir = dirs[0];
				i = 1;
			} else {
				i = 0;
			}
			
			for (int k = 1; k < count && i < number; ++k) {
				if (weights[k] > weights[0] * 0.8f || 
					fabs(InnerProduct(pSecondOrientation[index], dirs[k])) > 0.866f || 
					fabs(InnerProduct(pThirdOrientation[index], dirs[k])) > 0.866f) {
						(*seeds)[i].pos = pos;
						(*seeds)[i].dir = dirs[k];
						++i;
				}
			}

			if (pYMASK != NULL && pYMASK->GetScalar(pos) > 0.5f) {
				int y_id = -1;
				for (int k = 0; k < number; ++k) {
					if (fabs((*seeds)[k].dir.m_y) > (fabs((*seeds)[k].dir.m_x) + fabs((*seeds)[k].dir.m_z))) {
						y_id = k;
					}
				}
				if (y_id < 0) {
					;
				} else {
					Seed *temp_seeds = new Seed[number];
					memcpy(temp_seeds, *seeds, sizeof(Seed)*number);
					SafeDeleteArray(*seeds);
					*seeds = new Seed[number-1];
					for (int k = 0; k < y_id; ++k) {
						(*seeds)[k].pos = temp_seeds[k].pos;
						(*seeds)[k].dir = temp_seeds[k].dir;
					}
					for (int k = y_id+1; k < number; ++k) {
						(*seeds)[k-1].pos = temp_seeds[k].pos;
						(*seeds)[k-1].dir = temp_seeds[k].dir;
					}
					SafeDeleteArray(temp_seeds);
					number--;
				}
			}

			return number;
		}
	}
}

int GetNextDirection(const CPoint3F pos, CVector3F *in, CVector3F *out, const float threshold)
{
	int model = ModelSelection(pos);

	if (model == 0) {
		return 0;
	} 
	
	float filteredDTI[7];
	DTIMLSFilter(pos, filteredDTI);
	float fa = pDTI->DTI2FA(filteredDTI);
	//if (fa < 0.1f || pWMMASK->GetScalar(pos) > 0.4f)
	//	return 0;
	if (fa < 0.1f)
		return 0;

	if (model == 1) {
		/* use DTI to compute the direction */
		
		CVector3F dir;
		pDTI->DTI2Direction(filteredDTI, &dir);

		out->m_x = dir.m_x;
		out->m_y = dir.m_y;
		out->m_z = dir.m_z;

		float v = InnerProduct(in, out);
		if (fabs(v) > threshold) {
			if (v < 0.0f) {
				out->m_x = -out->m_x;
				out->m_y = -out->m_y;
				out->m_z = -out->m_z;
			}
			return 1;
		} else {
			;
			/* try HOT to find possiable direction */
		}
	} 

	/* try 4th order tensor first */

	/* use HOT to compute the direction */
	int n = pHOT->GetN();
	float *filteredHOT = new float[n];
	//HOTMLSFilter(pos, filteredHOT, filteredDTI);
	HOTMLSFilter(pos, filteredHOT);

	CVector3F dirs[3];
	int count = pHOT->HOT2Directions(filteredHOT, dirs);
	delete[] filteredHOT;

	float max_v = 0.0f;
	int index = 0;
	for (int i = 0; i < count; ++i) {
		float v = InnerProduct(in, &(dirs[i]));
		if (fabs(v) > fabs(max_v)) {
			max_v = v;
			index = i;
		}
	}

	out->m_x = dirs[index].m_x;
	out->m_y = dirs[index].m_y;
	out->m_z = dirs[index].m_z;

	float v = InnerProduct(in, out);
	if (fabs(v) > threshold) {
		if (v < 0.0f) {
			out->m_x = -out->m_x;
			out->m_y = -out->m_y;
			out->m_z = -out->m_z;
		}
		AddOrientation(pos, *out);
		return 1;
	} else {
		;		/* try 6th order tensor */
	}

	/* try 6th order tensor */
	n = pHOT6->GetN();
	float *filteredHOT6 = new float[n];
	//HOTMLSFilter(pos, filteredHOT, filteredDTI);
	HOT6MLSFilter(pos, filteredHOT6);

	count = pHOT6->HOT2Directions(filteredHOT6, dirs);
	delete[] filteredHOT6;

	max_v = 0.0f;
	index = 0;
	for (int i = 0; i < count; ++i) {
		float v = InnerProduct(in, &(dirs[i]));
		if (fabs(v) > fabs(max_v)) {
			max_v = v;
			index = i;
		}
	}

	out->m_x = dirs[index].m_x;
	out->m_y = dirs[index].m_y;
	out->m_z = dirs[index].m_z;

	v = InnerProduct(in, out);
	if (fabs(v) > threshold) {
		if (v < 0.0f) {
			out->m_x = -out->m_x;
			out->m_y = -out->m_y;
			out->m_z = -out->m_z;
		}
		AddOrientation(pos, *out);
		return 1;
	} else {
		return 0;		/* cannot find next direction */
	}
}

int ComputeNextPositionRK2(CPoint3F lastPos, CVector3F lastDir, 
	CPoint3F &nextPos, CVector3F &nextDir, const float threshold)
{
	CVector3F k1 = lastDir * (stepSize / 2.0f);
	CPoint3F p = lastPos + k1;
	CVector3F v;
	if (GetNextDirection(p, &lastDir, &v, threshold) == 0) {
		return 0;
	}

	CVector3F k2 = v * stepSize;

	nextPos = lastPos + k2;

	if (GetNextDirection(nextPos, &lastDir, &nextDir, threshold) == 0) {
		return 0;
	}

	return 1;
}

float Tracking(CFiber *fiber, Seed seed)
{
	float sum = fabs(seed.dir.m_x) + fabs(seed.dir.m_y) + fabs(seed.dir.m_z);
	if (sum < 0.5f) {
		fiber->CreateFiber(&(seed.pos), &(seed.pos), 1, 1, &seed);
		return 0.0f;
	}


	float initDTI[7] = {1, 1, 0, 0, 1, 0, 1};
	float filteredDTI[7];
	DTIMLSFilter(seed.pos, filteredDTI);

	int innerSteps = 2;
	uint maxSteps = (uint)(maxLength / (stepSize * (float)innerSteps));
	uint fnum, bnum;
	CPoint3F *pF = new CPoint3F[maxSteps];
	CPoint3F *pB = new CPoint3F[maxSteps];
	CPoint3F lastPos;
	CVector3F lastDir;

	/* tracking forward */
	lastPos = seed.pos;
	lastDir = seed.dir;
	pF[0] = lastPos;
	for (fnum = 1; fnum < maxSteps; ++fnum) {
		CPoint3F nextPos;
		CVector3F nextDir;

		bool flag = true;
		for (int step = 0; step < innerSteps; ++step) {
			if (ComputeNextPositionRK2(lastPos, lastDir, nextPos, nextDir, fiber->m_threshold) == 0) {
				flag = false;
				break;
			} 
			lastDir = nextDir;
			lastPos = nextPos;
		}
		if (flag == false)
			break;
		
		pF[fnum] = nextPos;
		lastDir = nextDir;
		lastPos = nextPos;
	}

	/* tracking backward */
	lastPos = seed.pos;
	lastDir.m_x = -seed.dir.m_x;
	lastDir.m_y = -seed.dir.m_y;
	lastDir.m_z = -seed.dir.m_z;
	pB[0] = lastPos;
	for (bnum = 1; bnum < maxSteps; ++bnum) {
		CPoint3F nextPos;
		CVector3F nextDir;

		bool flag = true;
		for (int step = 0; step < innerSteps; ++step) {
			if (ComputeNextPositionRK2(lastPos, lastDir, nextPos, nextDir, fiber->m_threshold) == 0) {
				flag = false;
				break;
			}
			lastDir = nextDir;
			lastPos = nextPos;
		}
		if (flag == false)
			break;
		
		pB[bnum] = nextPos;
		lastDir = nextDir;
		lastPos = nextPos;
	}

	fiber->CreateFiber(pF, pB, fnum, bnum, &seed);

	SafeDeleteArray(pF);
	SafeDeleteArray(pB);

	//return float(stepSize * (bnum + fnum - 2) * innerSteps);
	if (fnum > bnum) {
		return float(stepSize * (bnum - 1) * innerSteps);
	} else {
		return float(stepSize * (fnum - 1) * innerSteps);
	}
}

void ComputeSecondAndThirdOrientation()
{
	int sw = pDTI->GetWidth()* scale;
	int sh = pDTI->GetHeight() * scale;
	int sd = pDTI->GetDepth() * scale;

	omp_set_num_threads(16);

	for (int y = 590; y < 650; ++y) {
		//CFiber **temp_fibers = new CFiber*[sw*sd];
		int *temp_count = new int[sw*sd];
		Seed **temp_seeds = new Seed*[sw*sd];
		//memset(temp_fibers, 0, sizeof(CFiber*)*sw*sd);
		memset(temp_count, 0, sizeof(int)*sw*sd);
		memset(temp_seeds, 0, sizeof(Seed*)*sw*sd);

		for (int z = 0; z < sd; ++z) {
#pragma omp parallel for
			for (int x = 0; x < sw; ++x) {
				uint index = z * sw + x;
				CPoint3F pos = CPoint3F();
				pos.m_x = (float)x / (float)scale;
				pos.m_y = (float)y / (float)scale;
				pos.m_z = (float)z / (float)scale;
				temp_count[index] = GenerateSeedsNew(pos, &(temp_seeds[index]));
			}
		}

		for (int z = 0; z < sd; ++z) {
#pragma omp parallel for schedule(dynamic, 1)
			for (int x = 0; x < sw; ++x) {
				int index = z * sw + x;
				CFiber temp_fiber;
				//temp_fibers[index] = new CFiber[temp_count[index]];
				for (int j = 0; j < temp_count[index]; ++j) {
					Tracking(&temp_fiber, temp_seeds[index][j]);
				}
			}
			float ff = float(z+1) / float(sd) * 100.0f;
			printf("\rFiber tracking on slice y = %d ... %.3f%%", y, ff);
		}

#pragma omp parallel for
		for (int i = 0; i < sd*sw; ++i) {
			//SafeDeleteArray(temp_fibers[i]);
			SafeDeleteArray(temp_seeds[i]);
		}
		//SafeDeleteArray(temp_fibers);
		SafeDeleteArray(temp_count);
		SafeDeleteArray(temp_seeds);
	}
	printf("\n");
}

void FiberTracking()
{
	CreateMLSSamplePoints();
	CreateMatG();

	int w = pDTI->GetWidth();
	int h = pDTI->GetHeight();
	int d = pDTI->GetDepth();

	int sw = w * scale;
	int sh = h * scale;
	int sd = d * scale;
	int dy = endY - startY;

	pSecondOrientation = new CVector3F[sw*dy*sd];
	pThirdOrientation = new CVector3F[sw*dy*sd];
	memset(pSecondOrientation, 0, sizeof(CVector3F)*sw*dy*sd);
	memset(pThirdOrientation, 0, sizeof(CVector3F)*sw*dy*sd);

	FILE *fp_second = fopen("C:\\Ziang\\Dataset\\HumanConnectome\\792564_processed\\For_Paper\\B3000\\8X_2_SecondOrientation.raw", "rb");
	fread(pSecondOrientation, sizeof(CVector3F), sw*dy*sd, fp_second);
	fclose(fp_second);
	FILE *fp_third = fopen("C:\\Ziang\\Dataset\\HumanConnectome\\792564_processed\\For_Paper\\B3000\\8X_2_ThirdOrientation.raw", "rb");
	fread(pThirdOrientation, sizeof(CVector3F), sw*dy*sd, fp_third);
	fclose(fp_third);

	/* create second and third orientation */
	//isTrackingFlag = false;
	//isFirstFlag = true;
	//ComputeSecondAndThirdOrientation();
	//isFirstFlag = false;
	//ComputeSecondAndThirdOrientation();

	/* save results */
	//fp_second = fopen("C:\\Ziang\\Dataset\\HumanConnectome\\792564_processed\\For_Paper\\B3000\\8X_2_SecondOrientation.raw", "wb");
	//fwrite(pSecondOrientation, sizeof(CVector3F), sw*dy*sd, fp_second);
	//fclose(fp_second);
	//fp_third = fopen("C:\\Ziang\\Dataset\\HumanConnectome\\792564_processed\\For_Paper\\B3000\\8X_2_ThirdOrientation.raw", "wb");
	//fwrite(pThirdOrientation, sizeof(CVector3F), sw*dy*sd, fp_third);
	//fclose(fp_third);

	FILE *fp = fopen(pFiberPathname, "w");
	fprintf(fp, "%d %d %d %d %f %f\n", sw, sh, sd, scale, stepSize, maxLength);

	std::string pathname = pFiberPathname;
	size_t pos = pathname.find_last_of(SLASH);
	std::string path = pathname.substr(0, pos+1);

	omp_set_num_threads(16);

	isFirstFlag = true;
	isTrackingFlag = true;
	for (int y = 612; y < 613; ++y) {
		printf("Fiber tracking on slice y = %d ... ", y);

		CFiber **fibers = new CFiber*[sw*sd];
		int *count = new int[sw*sd];
		Seed **seeds = new Seed*[sw*sd];
		memset(fibers, 0, sizeof(CFiber*)*sw*sd);
		memset(count, 0, sizeof(int)*sw*sd);
		memset(seeds, 0, sizeof(Seed*)*sw*sd);

		for (int z = 571; z < 620; ++z) {
#pragma omp parallel for
			for (int x = 295; x < 339; ++x) {
				uint index = z * sw + x;
				CPoint3F pos = CPoint3F();
				pos.m_x = (float)x/(float)scale;
				pos.m_y = (float)y/(float)scale;
				pos.m_z = (float)z/(float)scale;
				count[index] = GenerateSeedsNew(pos, &(seeds[index]));
			}
		}
	
		for (int z = 0; z < sd; ++z) {
#pragma omp parallel for schedule(dynamic, 1)
			for (int x = 0; x < sw; ++x) {
				int index = z * sw + x;
				if (count[index] < 1)
					continue;

				fibers[index] = new CFiber[count[index]];
				CFiber *temp_fibers = new CFiber[count[index]];
				float *len = new float[count[index]];

				//for (int j = 0; j < count[index]; ++j) {
				//	if (count[index] != 1 && fabs(seeds[index][j].dir.m_y) > 0.7f)
				//		fibers[index][j].CreateFiber(&(seeds[index][j].pos), &(seeds[index][j].pos), 1, 1, &seeds[index][j]);
				//	else
				//		len[j] = Tracking(&(fibers[index][j]), seeds[index][j]);
				//}

				for (int j = 0; j < count[index]; ++j) {
					if (count[index] != 1 && fabs(seeds[index][j].dir.m_y) > 0.7f) {
						temp_fibers[j].CreateFiber(&(seeds[index][j].pos), &(seeds[index][j].pos), 1, 1, &seeds[index][j]);
						len[j] = 0.0f;
					} else {
						if (fabs(seeds[index][j].dir.m_z) > 0.8f) {
							temp_fibers[j].m_threshold = 0.5f;
						} else {
							temp_fibers[j].m_threshold = 0.94f;
						}
						len[j] = Tracking(&(temp_fibers[j]), seeds[index][j]);
					}
				}
				if (count[index] == 1) {
					fibers[index][0].CopyFiber(&(temp_fibers[0]));
				} else {
					int final_count = 0;
					for (int j = 0; j < count[index]; ++j) {
						if (len[j] > 10.0f) {
							fibers[index][final_count].CopyFiber(&(temp_fibers[j]));
							final_count++;
						}
					}
					if (final_count == 0) {
						fibers[index][0].CopyFiber(&(temp_fibers[0]));
						final_count = 1;
					}
					count[index] = final_count;
				}
				//	if (len[0] < 20.0f) {
				//		float max_len = len[0];
				//		int max_id = 0;
				//		for (int j = 1; j < count[index]; ++j) {
				//			if (len[j] > max_len) {
				//				max_len = len[j];
				//				max_id = j;
				//			}
				//		}

				//		if (max_id != 0) {
				//			/* swap fiber 0 with max_id */
				//			CFiber tt;
				//			tt.CopyFiber(&(temp_fibers[0]));
				//			temp_fibers[0].CopyFiber(&(temp_fibers[max_id]));
				//			temp_fibers[max_id].CopyFiber(&tt);
				//		}
				//	}
				//	fibers[index][0].CopyFiber(&(temp_fibers[0]));
				//	int final_count = 1;
				//	for (int j = 1; j < count[index]; ++j) {
				//		if (len[j] > 20.0f) {
				//			fibers[index][final_count].CopyFiber(&(temp_fibers[j]));
				//			final_count++;
				//		}
				//	}
				//	count[index] = final_count;
				//}

				SafeDeleteArray(temp_fibers);
				SafeDeleteArray(len);
			}
			float ff = float(z+1) / float(sd) * 100.0f;
			printf("\rFiber tracking on slice y = %d ... %.3f%%", y, ff);
		}

		char confPathname[PATHNAME_SIZE];
		char vertPathname[PATHNAME_SIZE];
		sprintf(confPathname, "%sslice_y=%08.04f.conf", path.c_str(), (float)y/(float)scale);
		sprintf(vertPathname, "%sslice_y=%08.04f.vert", path.c_str(), (float)y/(float)scale);
		fprintf(fp, "%s\t%s\n", confPathname, vertPathname);

		//SaveFibers(confPathname, vertPathname, fibers, count, sw, sd);

		SaveFibersAsVTK("C:\\Ziang\\Dataset\\Alex\\2X\\ForFiberTracking\\test.vtk", fibers, count, sw, sd);

#pragma omp parallel for
		for (int i = 0; i < sd*sw; ++i) {
			SafeDeleteArray(fibers[i]);
			SafeDeleteArray(seeds[i]);
		}
		SafeDeleteArray(fibers);
		SafeDeleteArray(count);
		SafeDeleteArray(seeds);

		printf("\rFiber tracking on slice y = %d ... 100.0%% Done.\n", y);
	}

	fclose(fp);

	SafeDeleteArray(pSecondOrientation);
	SafeDeleteArray(pThirdOrientation);
}

int main(int argc, char *argv[])
{
	if (ParseArguments(argc, argv) != 0) {
		PrintUsage();
		return 0;
	}

	pDTI = new CDTI();
	pDTI->OpenDTIFile(pDTIPathname);
	pHOT = new CHOT();
	pHOT->OpenHOTFile(pHOTPathname);
	pHOT6 = new CHOT();
	pHOT6->OpenHOTFile(pHOT6Pathname);
	pQBall = new CHOT();
	pQBall->OpenHOTFile(pQBallPathname);
	pBrainMASK = new CScalar();
	pBrainMASK->OpenScalarFile(pBrainMASKPathname);
	pWMMASK = new CScalar();
	pWMMASK->OpenScalarFile(pWMMASKPathname);
	
	if (pModelSelectionPathname != NULL) {
		pModelSelection = new CScalar();
		pModelSelection->OpenScalarFile(pModelSelectionPathname);
	}
	if (pSinglePathname) {
		pSingleMASK = new CScalar();
		pSingleMASK->OpenScalarFile(pSinglePathname);
	}
	if (pYMASKPathname) {
		pYMASK = new CScalar();
		pYMASK->OpenScalarFile(pYMASKPathname);
	}

	FiberTracking();

	SafeDelete(pDTI);
	SafeDelete(pHOT);
	SafeDelete(pHOT6);
	SafeDelete(pBrainMASK);
	SafeDelete(pWMMASK);
	SafeDelete(pModelSelection);
	SafeDelete(pSingleMASK);
	SafeDelete(pYMASK);
	SafeDelete(pQBall);

	return 0;
}
